package com.simple.llm.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.github.zuihou.base.R;
import com.plexpt.chatgpt.ChatGPT;
import com.plexpt.chatgpt.ChatGPTStream;
import com.plexpt.chatgpt.api.Api;
import com.plexpt.chatgpt.entity.billing.Usage;
import com.plexpt.chatgpt.entity.chat.*;
import com.plexpt.chatgpt.util.Proxys;
import com.simple.llm.component.OpenAiInterfaceLLMStreamListener;
import com.simple.llm.config.LLMProperties;
import com.simple.llm.domain.dto.LLMStreamResultDTO;
import com.simple.llm.domain.dto.NormalChatDTO;
import com.simple.llm.domain.dto.StreamChatDTO;
import com.simple.llm.domain.vo.LLMResponseVO;
import com.simple.llm.service.LLMStrategy;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import retrofit2.Retrofit;
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
import retrofit2.converter.jackson.JacksonConverterFactory;

import javax.servlet.http.HttpServletResponse;
import java.net.Proxy;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
public class OpenAiInterfaceLLMStrategyImpl implements LLMStrategy {
    private final LLMProperties.GptProperties gptProperties;
    private final ChatGPTStream chatGPTStream;
    private final ChatGPT chatGPT;
    private final Consumer<LLMStreamResultDTO> defaultStreamHandler = (o) -> {
        log.info("流式对话结束，结果:{}", o);
    };

    public OpenAiInterfaceLLMStrategyImpl(LLMProperties.GptProperties gptProperties, ObjectMapper objectMapper) {
        this.gptProperties = gptProperties;
        Proxy proxy = null;
        if (StringUtils.isNotBlank(gptProperties.getProxyIp())) {
            proxy = Proxys.http(gptProperties.getProxyIp(), gptProperties.getProxyPort());
        }
        chatGPTStream = ChatGPTStream.builder()
                .apiKeyList(gptProperties.getApiKeys())
                .timeout(gptProperties.getTimeout())
                .apiHost(gptProperties.getApiHost())
                .proxy(proxy)
                .build()
                .init();

        chatGPT = ChatGPT.builder()
                .apiKeyList(gptProperties.getApiKeys())
                .timeout(gptProperties.getTimeout())
                .apiHost(gptProperties.getApiHost())
                .proxy(proxy)
                .build()
                .init();
        ;
        chatGPT.setApiClient(new Retrofit.Builder()
                .baseUrl(chatGPT.getApiHost())
                .client(chatGPT.getOkHttpClient())
                .addCallAdapterFactory(RxJava2CallAdapterFactory.create())
                .addConverterFactory(JacksonConverterFactory.create(objectMapper))
                .build()
                .create(Api.class));
    }

    private ChatCompletion getChatCompletion(NormalChatDTO normalChatDTO) {
        List<Message> msgList = normalChatDTO.getMessages().stream().map(o -> {
            Message message = new Message();
            BeanUtils.copyProperties(o, message);
            return message;
        }).collect(Collectors.toList());

        int maxToken = 1000;
        if (normalChatDTO.getMaxTokens() != null) {
            maxToken = normalChatDTO.getMaxTokens();
            if (gptProperties.getMaxContextTokens() != null && maxToken > gptProperties.getMaxContextTokens()) {
                maxToken = gptProperties.getMaxContextTokens();
            }
        }

        ResponseFormat responseFormat = null;
        if (StringUtils.isNotBlank(normalChatDTO.getResponseFormat())) {
            responseFormat = ResponseFormat.builder().type(normalChatDTO.getResponseFormat()).build();
        }

        return ChatCompletion.builder()
                .model(gptProperties.getModelName())
                .messages(msgList)
                .maxTokens(maxToken)
                .temperature((double) (normalChatDTO.getTemperature() == null ? 0.8f : normalChatDTO.getTemperature()))
                .topP((double) (normalChatDTO.getTopP() == null ? 0.8f : normalChatDTO.getTopP()))
                .responseFormat(responseFormat)
                .build();
    }

    @Override
    public R<LLMResponseVO> chat(NormalChatDTO normalChatDTO) {
        ChatCompletionResponse chatCompletionResponse = chatGPT.chatCompletion(getChatCompletion(normalChatDTO));
        Message message = chatCompletionResponse.getChoices().get(0).getMessage();
        Usage usage = chatCompletionResponse.getUsage();
        LLMResponseVO vo = LLMResponseVO.builder()
                .content(message.getContent())
                .inputTokenCount(usage.getPromptTokens())
                .outputTokenCount(usage.getCompletionTokens())
                .desc(normalChatDTO.getDesc())
                .createUser(normalChatDTO.getCreateUser())
                .build();
        return R.success(vo);
    }

    @Override
    public ResponseBodyEmitter streamChat(StreamChatDTO streamChatDTO, HttpServletResponse response, Consumer<LLMStreamResultDTO> resultHandler) {
        response.setContentType("application/octet-stream");
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(gptProperties.getStreamTimeOut() == null ? 60 * 1000L : gptProperties.getStreamTimeOut());
        OpenAiInterfaceLLMStreamListener openAiInterfaceLLMStreamListener = new OpenAiInterfaceLLMStreamListener(emitter, streamChatDTO, resultHandler == null ? defaultStreamHandler : resultHandler);
        ChatCompletion chatCompletion = getChatCompletion(streamChatDTO);
        chatCompletion.setStreamOptions(new StreamOption(true));
        chatGPTStream.streamChatCompletion(chatCompletion, openAiInterfaceLLMStreamListener);
        return emitter;
    }

}
